#include "AnaDef.h"
#include "XGLUtils.h"

void doContaminationFit(const Int_t itype, const Int_t ipar, const Int_t imom, const Bool_t kinvq, const TString cvar, const Double_t var0, const Double_t var1, const Double_t vard)
{
  //--------------------------------------------------------------------
  //                          Prelude
  //--------------------------------------------------------------------
  
  gErrorIgnoreLevel = 1001;
  gStyle->SetOptStat(0);
  //gPad->SetLogy();

  TTree *tree=(TTree*)gDirectory->Get("newdraw_Tree");
  if(!tree){
    gDirectory->ls();
    printf("no tree!\n");exit(1);
  }

  TList *ll=new TList;

  TGraphErrors *gr=new TGraphErrors;
  gr->SetName(cvar);
  ll->Add(gr);

  Double_t mass = -999;
  if(ipar==0){
    mass = 5.10998885147273540e-04;
  }
  else if(ipar==4){
    mass = 9.38271999359130859e-01;
  }
  else{
    printf("wrong ipar!! %d\n", ipar); exit(1);
  }

  tree->SetAlias("nch", Form("trdnch%d", kinvq));
  tree->SetAlias("rtpc","tpcsig/AliTRDdEdxUtils::QMeanTPC(pow(10,ltbg))");
  tree->SetAlias("trdT","trackTmeanV2");

  //pid==2, tpcsig*-1;
  //pid!=2, tpcsig*1;
  tree->SetAlias("pidtpc", "tpcsig*(-1+2*(pid!=2))");
  
  //--------------------------------------------------------------------
  //                          Define basecut
  //--------------------------------------------------------------------
  TString basecut;
  if(itype==2){
    tree->SetAlias("rELtrd",  Form("trdQ%d/AliTRDdEdxUtils::Q%dMeanTRDpp(mom/%e)", kinvq, kinvq, mass));

    basecut = Form("itype== %d ", itype);

    basecut += "&& run >= 126461 && run <= 130930 "; printf("selecting 10eeeeeeeeee !!\n");
    //basecut += "&& run >= 121527 && run <= 126460 "; printf("selecting 10dddddddddd !!\n");

    basecut += "&& rtpc>0.95 && rtpc<1.05 ";
  }
  else if(itype==4){
    tree->SetAlias("rELtrd",  Form("trdQ%d/AliTRDdEdxUtils::Q%dMeanTRDPbPb(mom/%e)", kinvq, kinvq, mass));
    basecut = "itype >=10 && itype <=11 ";

    basecut += "&& rtpc>0.95 && rtpc<1.3 ";
  }
  else{
    printf("error itype ! %d\n", itype); exit(1);
  }
  
  basecut += Form("&& (pid==%d || pid==2) ", ipar);

  basecut += "&& ifill==0 && nch>=6 ";

  TString pcut;
  if(imom==0){
    pcut="&& mom>0.5 && mom<1.0 ";
  }
  else if(imom==1){
    pcut="&& mom>1.0 && mom<2.0 ";
  }
  else if(imom==2){
    pcut="&& mom>2.0 && mom<5.0 ";
  }
  else{
    printf("imom error ! %d\n", imom);exit(1);
  }

  basecut+= pcut;
  
  printf("basecut: %s\n", basecut.Data());

  //--------------------------------------------------------------------
  //                             Get momentum
  //--------------------------------------------------------------------
  TH2D *hmom = new TH2D("hmom","",200,0,5, 5,0,5);
  tree->Draw("pid:mom>>hmom",basecut);
  const Double_t pmean = hmom->GetMean(1);
  const Double_t prms = hmom->GetRMS(1);
  const TString pinfo(Form("p = %.2f#pm%.2f GeV/c", pmean, prms));
  ll->Add(hmom);
  //--------------------------------------------------------------------
  //                            Before Cut
  //--------------------------------------------------------------------
  const Int_t nbin = 2000;
  const Double_t xmin = -200;
  const Double_t xmax =  200;
  const Int_t midbin = nbin/2;

  TH1D *hh=new TH1D("hh",basecut, nbin, xmin, xmax);
  tree->Draw(Form("pidtpc>>%s",hh->GetName()),basecut);
  ll->Add(hh);

  if(hh->GetBinContent(0) || hh->GetBinContent(nbin+1)){
    printf("hh overflow!! %f %f\n", hh->GetBinContent(0), hh->GetBinContent(nbin+1));
    TFile *fout=new TFile("ferror.root","recreate");
    hh->Write();
    fout->Save();
    fout->Close();
    exit(1);
  }

  const Double_t oPi=hh->Integral(0,midbin);
  const Double_t oEl=hh->Integral(midbin, nbin+1);
  printf("\n\nbefore pi: %15f e: %15f\n\n", oPi, oEl);
  
  //--------------------------------------------------------------------
  //                            After Cut
  //--------------------------------------------------------------------
  Double_t rcut = var0;
  Int_t npoint = 0;
  while(rcut<var1){

    printf("\n\n===============================================================\n");
    const TString pidcut(Form("&& %s > %f", cvar.Data(), rcut));

    printf("\n\ncut: %s pidcut: %s\n\n", basecut.Data(), pidcut.Data());
    
    TH1D *hc=new TH1D(Form("hc%d",npoint),basecut+pidcut, nbin, xmin, xmax);
    tree->Draw(Form("pidtpc>>%s",hc->GetName()),basecut + pidcut);
    ll->Add(hc);

    const Double_t nPi=hc->Integral(0,midbin);
    const Double_t nEl=hc->Integral(midbin, nbin+1);
  
    printf("\n\nafter pi: %15f e: %15f\n\n", nPi, nEl);
    
    if(nPi<50){
      printf("nPi<50! %.1f\n", nPi); break;
    }

    const Double_t Pieff = nPi/oPi;
    const Double_t Pierr = sqrt(Pieff*(1-Pieff)/oPi);
    const Double_t Eleff = nEl/oEl;
    const Double_t Elerr = sqrt(Eleff*(1-Eleff)/oEl);

    printf("%s %s npoint %d pidcut: %s before: %6.0f %6.0f after: %6.0f %6.0f PIeff: %6.6f %6.6f Eeff: %6.4f %6.4f\n", pcut.Data(), pinfo.Data(), npoint, pidcut.Data(), oPi, oEl, nPi, nEl, Pieff, Pierr, Eleff, Elerr);
    
    gr->SetPoint(npoint, Eleff, Pieff);
    gr->SetPointError(npoint, Elerr, Pierr);
    npoint++;

    rcut += vard;
  }

  TCanvas *cc=new TCanvas("cc");
  gr->SetTitle(pinfo);
  gr->SetMarkerStyle(20);
  gr->SetMarkerSize(1);
  gr->GetXaxis()->SetLimits(0.5,1);
  gr->SetMinimum(5e-3);
  gr->SetMaximum(1);
  gr->Draw("ap");
  gPad->SetGrid();

  TFile *fout=new TFile(Form("scanFit_par%d_mom%d_type%d_Q%d_var%s.root", ipar, imom, itype, kinvq, cvar.Data()),"recreate");
  ll->Write();
  cc->Write();
  fout->Save();
  fout->Close();
}

int main(const int argc, char *argv[])
{
  int argi[argc];
  double argf[argc];
  for(Int_t ii=0; ii<argc; ii++){
    printf("%d: %s\n", ii, argv[ii]);
    argi[ii] = atoi(argv[ii]);
    argf[ii] = atof(argv[ii]);
  }

  if(argc!=10){
    printf("argc!=\n"); exit(1);
  }

  TFile::Open(argv[1]);

  const Int_t itype   = argi[2];
  const Bool_t kinvq  = argi[3];
  const TString var   = argv[4];
  const Double_t cut0 = argf[5];
  const Double_t cut1 = argf[6];
  const Double_t dcut = argf[7];
  const Int_t imom    = argi[8];
  const Int_t ipar    = argi[9];

  doContaminationFit(itype, ipar, imom, kinvq, var, cut0, cut1, dcut);

  return 0;
}

  
